import gymnasium as gym
from stable_baselines3 import PPO
import torch

# 自动选择设备（MLP 用 CPU 更好）
device = "cpu"
print("Using device:", device)

# 创建环境
env = gym.make("CartPole-v1", render_mode="human")

# 创建 PPO 模型
model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    tensorboard_log="./cartpole_log/",
    device=device
)

# 训练 100k 步
print("开始训练 CartPole...")
model.learn(total_timesteps=100000)
print("训练结束 ✅")

# 保存模型
model.save("ppo_cartpole")

# 测试模型
obs, _ = env.reset()
for _ in range(2000):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, _ = env.step(action)
    if terminated or truncated:
        obs, _ = env.reset()

env.close()

